import hashlib
from contextlib import contextmanager
from pysqlcipher3 import dbapi2 as sqlite3

# rotki-pysqlcipher3==2024.10.1

class V4Decryptor:
    """
    V4 版本特定常量
    """
    PageSize = 4096
    V4IterCount = 256000
    HmacSHA512Size = 64

    KeySize = 32
    SaltSize = 16
    AESBlockSize = 16
    SQLiteHeader = "SQLite format 3\x00"
    IVSize = 16

    def MakeCipherKey(self, db_path, key):
        """
        扩展出微信数据库的真实密钥
        """
        with open(db_path, "rb") as f:
            mac_salt = f.read(self.SaltSize)
            rawKey = bytes.fromhex(key)
            mac_key = hashlib.pbkdf2_hmac('sha512', rawKey, mac_salt, self.V4IterCount, self.KeySize)
            realKey = (mac_key + mac_salt).hex()
            return realKey

    @contextmanager
    def OpenDb(self, db_path, password, read_only=True):
        # 扩展出真实的密钥
        real_key = self.MakeCipherKey(db_path, password)

        if read_only:
            # 使用URI模式并添加mode=ro参数
            uri = f"file:{db_path}?mode=ro"
            _db = sqlite3.connect(uri, uri=True)
        else:
            _db = sqlite3.connect(db_path)

        print(f"db_path={db_path}, real_key={real_key}")

        try:
            _db.execute("PRAGMA wal_autocheckpoint = 0;")
            _db.execute("PRAGMA checkpoint_fullfsync = true;")
            _db.execute("PRAGMA temp_store = 1;")
            _db.execute("PRAGMA key = \"x'{}'\";".format(real_key))
            # 这两步操作必须在set database key之后
            _db.execute("PRAGMA cipher_page_size = 4096;")
            _db.execute("PRAGMA journal_mode = \"WAL\";")
            # 返回数据库连接对象
            yield _db
        finally:
            # 确保无论是否发生异常，连接都会被关闭
            _db.close()


if __name__ == "__main__":
    ddd = V4Decryptor()
    db_path = r""
    key = ""
    ppp = ddd.MakeCipherKey(db_path=db_path, key=key)
    print(ppp)

    ret_list = []
    with ddd.OpenDb(db_path=db_path, password=key, read_only=True) as db:
        sql = "SELECT username, alias, nick_name, remark FROM contact LIMIT 10"
        rows = db.execute(sql).fetchall()
        for row in rows:
            (username, alias, nick_name, remark) = row
            ret_list.append({
                "username": username,  # 微信id
                "alias": alias,  # 微信号
                "nick_name": nick_name,  # 昵称
                "remark": remark,  # 备注
            })
    print(ret_list)
